import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import time
from scipy.stats import linregress

plt.rcParams["font.sans-serif"] = ["SimHei"] # 设置显示中文字体 宋体
plt.rcParams["axes.unicode_minus"] = False

# ==================== 初始化设置 ====================
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')  # 替换为实际token
pro = ts.pro_api()

# 回测参数
initial_capital = 1_000_000  # 初始资金
start_date = '20190101'      # 回测开始日期
end_date = '20221231'        # 回测结束日期
stock_pool_size = 20         # 股票池数量
commission = 0.0003          # 交易佣金
slippage = 0.001            # 滑点系数

# 策略参数
short_ma = 5                # 短期均线周期
long_ma = 20                # 长期均线周期
stop_loss = 0.10            # 止损比例
take_profit = 0.20          # 止盈比例

# ==================== 数据准备模块 ====================
def fetch_stock_data():
    """获取并预处理股票数据"""
    print("正在获取成分股数据...")
    df_index = pro.index_weight(index_code='399300.SZ',
                                start_date=start_date,
                                end_date=end_date)
    stock_list = df_index['con_code'].unique().tolist()[:stock_pool_size]

    all_data = {}
    for ts_code in tqdm(stock_list, desc="下载股票数据"):
        for _ in range(3):  # 重试机制
            try:
                df = pro.daily(ts_code=ts_code,
                               start_date=start_date,
                               end_date=end_date)
                if not df.empty:
                    # 数据预处理
                    df = df.sort_values('trade_date')
                    df['trade_date'] = pd.to_datetime(df['trade_date'])
                    df.set_index('trade_date', inplace=True)
                    # 计算技术指标
                    df['short_ma'] = df['close'].rolling(short_ma).mean()
                    df['long_ma'] = df['close'].rolling(long_ma).mean()
                    df = df[df['long_ma'].notna()]
                    all_data[ts_code] = df
                break
            except Exception as e:
                print(f"下载{ts_code}数据失败: {str(e)}")
                time.sleep(1)
    return all_data


# ==================== 信号生成模块 ====================
def generate_signals(df):
    """生成交易信号"""
    df = df.copy()
    # 生成原始信号
    df['position'] = 0
    golden_cross = (df['short_ma'] > df['long_ma']) & (df['short_ma'].shift(1) <= df['long_ma'].shift(1))
    death_cross = (df['short_ma'] < df['long_ma']) & (df['short_ma'].shift(1) >= df['long_ma'].shift(1))
    df.loc[golden_cross, 'position'] = 1
    df.loc[death_cross, 'position'] = -1

    # 信号处理（兼容未来版本）
    df['position'] = (
        df['position']
        .replace(0, np.nan)
        .ffill()
        .fillna(0)
    )
    return df


# ==================== 风险管理模块 ====================
class RiskManager:
    def __init__(self, stop_loss, take_profit):
        self.stop_loss = stop_loss
        self.take_profit = take_profit
        self.positions = {}  # {ts_code: {'entry_price':, 'highest_price':}}

    def check_risk(self, ts_code, current_price):
        """检查止损止盈"""
        if ts_code not in self.positions:
            return False

        entry = self.positions[ts_code]['entry_price']
        highest = self.positions[ts_code]['highest_price']

        current_drawdown = (highest - current_price) / highest
        current_profit = (current_price - entry) / entry

        if current_drawdown >= self.stop_loss:
            return 'stop_loss'
        elif current_profit >= self.take_profit:
            return 'take_profit'
        return False

    def update_position(self, ts_code, price):
        """更新最高价"""
        if ts_code in self.positions:
            self.positions[ts_code]['highest_price'] = max(
                self.positions[ts_code]['highest_price'], price)


# ==================== 回测引擎 ====================
def backtest_strategy(data_dict):
    """增强版回测引擎"""
    cash = initial_capital
    portfolio_value = [cash]
    cash_history = [cash]
    holdings_history = [0]
    returns_history = [0.0]
    position_dict = {}
    risk_manager = RiskManager(stop_loss, take_profit)
    trade_log = []

    dates = pd.date_range(start=start_date, end=end_date)

    for date in tqdm(dates, desc="回测进度"):
        current_date = date.strftime('%Y%m%d')
        daily_holdings = 0

        # 平仓处理
        for ts_code in list(position_dict.keys()):
            try:
                df = data_dict[ts_code]
                if date in df.index:
                    close_price = df.loc[date, 'close']
                    risk_status = risk_manager.check_risk(ts_code, close_price)
                    if risk_status:
                        shares = position_dict.pop(ts_code)
                        close_value = shares * close_price
                        cash += close_value * (1 - commission - slippage)
                        risk_manager.positions.pop(ts_code)
                        trade_log.append({
                            'date': date,
                            'code': ts_code,
                            'action': 'sell',
                            'price': close_price,
                            'shares': shares,
                            'reason': risk_status
                        })
            except KeyError:
                continue

        # 开仓处理
        for ts_code, df in data_dict.items():
            if date in df.index:
                signal = df.loc[date, 'position']
                close_price = df.loc[date, 'close']

                if ts_code in position_dict:
                    risk_manager.update_position(ts_code, close_price)
                    daily_holdings += position_dict[ts_code] * close_price
                    continue

                if signal == 1 and cash > 10000:
                    position_size = min(cash * 0.1, cash - 10000)
                    shares = int(position_size // (close_price * (1 + commission + slippage)))
                    if shares > 0:
                        position_dict[ts_code] = shares
                    invest = shares * close_price * (1 + commission + slippage)
                    cash -= invest
                    risk_manager.positions[ts_code] = {
                        'entry_price': close_price,
                        'highest_price': close_price
                    }
                    trade_log.append({
                        'date': date,
                        'code': ts_code,
                        'action': 'buy',
                        'price': close_price,
                        'shares': shares,
                        'investment': invest
                    })

                    # 计算当日资产
                    total_assets = cash + daily_holdings
                    portfolio_value.append(total_assets)
                    cash_history.append(cash)
                    holdings_history.append(daily_holdings)
                    current_return = (total_assets / initial_capital - 1) * 100
                    returns_history.append(current_return)

                    # 定期显示
                    if len(portfolio_value) % 50 == 0:
                        print(f"\n日期：{date.strftime('%Y-%m-%d')}")
                    print(f"剩余资金：{cash:,.2f} | 持仓：{daily_holdings:,.2f}")
                    print(f"总资产：{total_assets:,.2f} | 收益率：{current_return:.2f}%")

    return portfolio_value, cash_history, holdings_history, returns_history, trade_log


# ==================== 绩效分析模块 ====================
def analyze_performance(results):
    """综合绩效分析"""
    portfolio_values, cash_history, holdings_history, returns_history, trades = results

    analysis_df = pd.DataFrame({
        'Total': portfolio_values[1:],
        'Cash': cash_history[1:],
        'Holdings': holdings_history[1:],
        'Return%': returns_history[1:]
    }, index=pd.date_range(start=start_date, periods=len(portfolio_values) - 1))

    # 计算指标
    total_return = (portfolio_values[-1] / initial_capital - 1) * 100
    max_drawdown = (analysis_df['Total'].cummax() - analysis_df['Total']).max()
    volatility = analysis_df['Total'].pct_change().std() * np.sqrt(252)

    print("\n======== 最终绩效报告 ========")
    print(f"初始资金: {initial_capital:,.2f}")
    print(f"最终资金: {portfolio_values[-1]:,.2f}")
    print(f"累计收益率: {total_return:.2f}%")
    print(f"年化波动率: {volatility:.2%}")
    print(f"最大回撤: {max_drawdown:,.2f}")
    print(f"交易次数: {len(trades)}")

    # 可视化
    plt.figure(figsize=(14, 10))

    # 资金曲线
    plt.subplot(3, 1, 1)
    analysis_df['Total'].plot(label='总资产', color='navy')
    analysis_df['Cash'].plot(label='现金', linestyle='--', color='green')
    analysis_df['Holdings'].plot(label='持仓', linestyle=':', color='orange')
    plt.title('资金分布变化')
    plt.ylabel('金额')
    plt.legend()
    plt.grid(True)

    # 收益率曲线
    plt.subplot(3, 1, 2)
    analysis_df['Return%'].plot(color='darkred')
    plt.title('累计收益率')
    plt.ylabel('百分比')
    plt.grid(True)

    # 每日收益分布
    plt.subplot(3, 1, 3)
    daily_returns = analysis_df['Total'].pct_change().dropna()
    daily_returns.hist(bins=50, color='teal', alpha=0.7)
    plt.title('每日收益分布')
    plt.xlabel('收益率')
    plt.grid(True)

    plt.tight_layout()
    plt.show()


# ==================== 主程序 ====================
if __name__ == "__main__":
    # 数据准备
    print("正在准备数据...")
    raw_data = fetch_stock_data()
    processed_data = {code: generate_signals(df) for code, df in raw_data.items()}

    # 执行回测
    print("\n开始回测...")
    results = backtest_strategy(processed_data)

    # 绩效分析
    analyze_performance(results)

    # 显示交易记录示例
    print("\n前5笔交易记录：")
    print(pd.DataFrame(results[-1][:5]).to_string(index=False))